import os
import shutil
import random


def labelGO_count(directory):
    """
    计算通过 labelGO 自动标注后的 txt文件和 jpg文件个数
    """
    # Iterate through all files
    all_files = os.listdir(directory)
    # Take out all files whose file extension is .txt
    txt_list = [os.path.splitext(file)[0] for file in all_files if os.path.splitext(file)[1] == '.txt']
    print('Number of txt files: ', len(txt_list))
    jpg_list = [os.path.splitext(file)[0] for file in all_files if os.path.splitext(file)[1] == '.jpg']
    print('Number of jpg files: ', len(jpg_list))
    return txt_list


def labelGO_copy(root_path):
    """
    以txt文件个数为标准，将labelGO自动标注的 txt文件和jpg文件移动到对应位置
    :param root_path:
    :return:
    """
    # Create txt files storage directory
    if not os.path.exists(os.path.join(root_path, 'Annotations')):
        os.mkdir(os.path.join(root_path, 'Annotations'))
    else:
        print(os.path.join(root_path, 'Annotations') + 'has exists!')
    # Create jpg files storage directory
    if not os.path.exists(os.path.join(root_path, 'Images_comple')):
        os.mkdir(os.path.join(root_path, 'Images_comple'))
    else:
        print(os.path.join(root_path, 'Images_comple') + ' has exists!')
    # Create all files storage directory
    if not os.path.exists(os.path.join(root_path, 'Storage')):
        os.mkdir(os.path.join(root_path, 'Storage'))
    else:
        print(os.path.join(root_path, 'Storage') + ' has exists!')
    # Delete duplicate jpg files and restore files in new directories
    classes_txt = os.path.join(root_path, 'classes.txt')
    shutil.move(classes_txt, os.path.join(root_path, 'Storage', 'classes.txt'))
    txt_files = labelGO_count(root_path)
    for txt in txt_files:
        # Get complete path of files
        origin_jpg = os.path.join(root_path, txt + '.jpg')
        origin_txt = os.path.join(root_path, txt + '.txt')
        jpg_des_path = os.path.join(root_path, 'Images_comple', txt + '.jpg')
        # txt_des_path = os.path.join(root_path, 'Annotations', txt + '.txt')
        jpg_storage_path = os.path.join(root_path, 'Storage', txt + '.jpg')
        txt_storage_path = os.path.join(root_path, 'Storage', txt + '.txt')
        # Copy files to designated directory
        shutil.copyfile(origin_jpg, jpg_des_path)
        shutil.copyfile(origin_jpg, jpg_storage_path)
        # shutil.copyfile(origin_txt, txt_des_path)
        shutil.copyfile(origin_txt, txt_storage_path)


def divide_dataset(root_path):
    # root_path = r'D:\Project\ShandongMine\DownMing_Img'
    Img_path = os.path.join(root_path, 'Images_comple')
    train_txt_path = os.path.join(root_path, 'train.txt')
    val_txt_path = os.path.join(root_path, 'val.txt')
    train_weight = 0.9
    Img_list = os.listdir(Img_path)
    train_set = random.sample(Img_list, int(len(Img_list) * train_weight))
    val_set = list(set(Img_list).difference(set(train_set)))
    with open(train_txt_path, 'w') as f1, open(val_txt_path, 'w') as f2:
        for i in train_set:
            f1.write(i + '\n')
        for j in val_set:
            f2.write(j + '\n')


def add_content(root_path, content):
    train_txt = os.path.join(root_path, 'train.txt')
    val_txt = os.path.join(root_path, 'val.txt')
    with open(train_txt, 'r') as f1:
        old_content = f1.readlines()
    with open(train_txt, 'w') as f1:
        for i in old_content:
            f1.write(os.path.join(content, i))
    with open(val_txt, 'r') as f2:
        old_content = f2.readlines()
    with open(val_txt, 'w') as f2:
        for j in old_content:
            f2.write(os.path.join(content, i))


if __name__ == '__main__':
    # Number of files
    root_path = r'/home/stark/algo-env/datasets/storage'
    labelGO_copy(root_path)
    # add_content(root_path, r'.\DownMing_Img\Images_comple')
